package leetcode.design;

import java.util.*;

class Solution398 {

    private final Map<Integer, List<Integer>> indexMap;
    private Random random;

    public Solution398(int[] nums) {
        indexMap = new HashMap<>();
        random = new Random();
        init(nums);
    }

    private void init(int[] nums) {
        for (int i = 0; i < nums.length; i++) {
            if (indexMap.get(nums[i]) == null) {
                List<Integer> list = new ArrayList<>();
                list.add(i);
                indexMap.put(nums[i], list);
            } else {
                indexMap.get(nums[i]).add(i);
            }
        }
    }

    public int pick(int target) {
        List<Integer> list = indexMap.get(target);
        int bound = random.nextInt(list.size());
        return list.get(bound);
    }
}

/**
 * Your Solution object will be instantiated and called as such:
 * Solution obj = new Solution(nums);
 * int param_1 = obj.pick(target);
 */

public class leetcode398 {
    public static void main(String[] args) {
        int[] nums = {1, 2, 3, 3, 3};
        Solution398 solution = new Solution398(nums);
        System.out.println(solution.pick(3)); // 随机返回索引 2, 3 或者 4 之一。每个索引的返回概率应该相等。
        System.out.println(solution.pick(1)); // 返回 0 。因为只有 nums[0] 等于 1 。
        System.out.println(solution.pick(3)); // 随机返回索引 2, 3 或者 4 之一。每个索引的返回概率应该相等。
    }
}
